import os
from os.path import join
import numpy as np
import open3d as o3d


def get_visualization_meshes(root_dir, scene_names, color=[0.0, 0.0, 0.0]):
    meshes = []
    for scene_name in scene_names:
        scene_motion_dir = join(root_dir, scene_name, "h1_kinematic_motions")
        for motion_filename in os.listdir(scene_motion_dir):
            data = np.load(join(scene_motion_dir, motion_filename), allow_pickle=True)["arr_0"].item()
            root_global_t = data["global_translations"][::1]  # (N//10, 3)
            lines, points = [], [root_global_t[0]]
            for i in range(1, root_global_t.shape[0]):
                lines.append([i-1, i])
                points.append(root_global_t[i])
            line_set = o3d.geometry.LineSet()
            line_set.lines = o3d.utility.Vector2iVector(np.int32(lines))
            line_set.colors = o3d.utility.Vector3dVector(np.float32(color).reshape(1, 3).repeat(len(lines), axis=0))
            line_set.points = o3d.utility.Vector3dVector(np.float32(points))
            meshes.append(line_set)

    return meshes


if __name__ == "__main__":
    training_scene_names = ['1006be65e7bc937e9141f9b58470d646', '1007e20d5e811b308351982a6e40cf41', '100b18376b885f206ae9ad7e32c4139d', '1013f70851210a618f2e765c4a8ed3d', '1015e71a0d21b127de03ab2a27ba7531', '1016f4debe988507589aae130c1f06fb', '1022fe7dd03f6a4d4d5ad9f13ac9f4e7', '1028b32dc1873c2afe26a3ac360dbd4', '1031fc859dc3177a2f84cb7932f866fd', '1033ee86cc8bac4390962e4fb7072b86', '103a0a413d4c3353a723872ad91e4ed1', '103a60f3b09107df2da1314e036b435e', '103c31671f8c0b1467bb14b25f99796e', '103d77d63f0d68a044e6721e9db29c1b', '1049953406c81b237eaeab1f0c9120b7', '1055f78d441d170c4f3443b22038d340', '1063d4fcd366de4060e37b3f76995f8b', '106a0dbaead5066519fb4103277a6b93', '106c7f10c5bf5bd5f51f77a6d7299806', '10709332176024ce9e47e7a22e24daa3', '1079635b3da12a812cee4bf5d0f11ffe', '107caefdad02cf1c8ab8e68cb52baa6a', '107ed94869ed6f1be13496cd332ce78f', '108238b535eb293cd79b19c7c4f0e293']
    test_scene_names = ['104256e5bb73b0b719fb4103277a6b93', '106e9487a1d47ff1a09cb462b3e723e2', '108b9cb292fd811cf51f77a6d7299806', '1093d35c2ac73bb74ca84d60642ec7e8', '10991b3b01cc6cc99982a4d6320875e']
    root_dir = "/home/liuyun/Humanoid_IL_Benchmark/retargeting/ROAM_retargeted_data"

    training_meshes = get_visualization_meshes(root_dir, training_scene_names, color=[1.0, 0.0, 0.0])
    test_meshes = get_visualization_meshes(root_dir, test_scene_names, color=[0.0, 1.0, 0.0])

    o3d.visualization.draw_geometries(training_meshes + test_meshes)

